from langchain_core.prompts import FewShotPromptTemplate, PromptTemplate
from langchain_openai import ChatOpenAI
from globle import envs
from langchain_core.output_parsers import StrOutputParser

# 用中文写一个Few-shot Prompt
example_prompt = PromptTemplate(
    input_variables=["input", "output"],
    template="Question: {input}\nAnswer: {output}",
)
examples = [
    {"input": "Who was the father of Mary Ball Washington?", "output": "George Washington"},
    {"input": "Who wrote the book '1984'?", "output": "George Orwell"},
]

prompt = FewShotPromptTemplate(
    prefix="请按照示例回答问题",
    examples=examples, # 传示例
    example_prompt=example_prompt, # 示例模板
    suffix="Question: {input}", # 提示后缀
    input_variables=["input"], # 输入变量
)

# 创建模型
model = ChatOpenAI(
    model="deepseek-chat",  # 或具体模型名
    base_url="https://api.deepseek.com/v1",  # DeepSeek API 地址
)


chain = prompt|model|StrOutputParser()
r = chain.invoke({'input':'谁是Mary Ball Washington的父亲?'})
print(r)